import pymysql
import pandas as pd


class DBManage:
    def __init__(self):
        self.connection = None
        self.cursor = None
        self.connect()

    def connect(self):
        """建立数据库连接"""
        try:
            self.connection = pymysql.connect(
                host='',
                user='',
                password='',
                db='',
                port=3306
            )
            self.cursor = self.connection.cursor()
        except pymysql.MySQLError as e:
            raise Exception(f"数据库连接失败: {e}")

    def __enter__(self):
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        self.close()

    def get_columns(self, table):
        """获取表字段"""
        self.refresh_connection()
        try:
            self.cursor.execute(f"DESC {table}")
            columns = self.cursor.fetchall()
            return [column[0] for column in columns]
        except pymysql.MySQLError as e:
            raise Exception(f"获取字段失败: {e}")

    def show_data(self, table):
        """显示当前表的数据"""
        columns = self.get_columns(table)
        if columns:
            sql = f"SELECT * FROM {table}"
            df = pd.DataFrame(self.execute_query(sql), columns=columns)
            return df.to_string(index=False)
        else:
            raise Exception("表不存在!")

    def add_data(self, table, data: dict):
        """添加数据"""
        self.refresh_connection()
        keys = ', '.join(data.keys())
        placeholders = ', '.join(['%s'] * len(data))
        sql = f"INSERT INTO {table} ({keys}) VALUES ({placeholders})"
        self.execute_update(sql, list(data.values()))

    def select_data(self, table, condition=None):
        """查询数据"""
        self.refresh_connection()
        sql = f"SELECT * FROM {table}"
        if condition:
            sql += f" WHERE {condition}"
        return self.execute_query(sql)

    def delete_data(self, table, condition):
        """删除数据"""
        self.refresh_connection()
        sql = f"DELETE FROM {table} WHERE {condition}"
        self.execute_update(sql)

    def update_data(self, table, updates: dict, condition):
        """更改数据"""
        self.refresh_connection()
        update_str = ', '.join([f"{col} = %s" for col in updates.keys()])
        sql = f"UPDATE {table} SET {update_str} WHERE {condition}"
        self.execute_update(sql, list(updates.values()))

    def execute_query(self, sql, params=None):
        """执行查询并返回结果"""
        try:
            self.cursor.execute(sql, params or ())
            return self.cursor.fetchall()
        except pymysql.MySQLError as e:
            raise Exception(f"查询执行失败: {e}")

    def execute_update(self, sql, params=None):
        """执行更新操作"""
        try:
            self.cursor.execute(sql, params or ())
            self.connection.commit()
        except pymysql.MySQLError as e:
            self.connection.rollback()
            raise Exception(f"数据更新失败: {e}")

    def close(self):
        """关闭数据库连接"""
        if self.cursor:
            self.cursor.close()
        if self.connection:
            self.connection.close()

    def refresh_connection(self):
        """刷新数据库连接"""
        self.close()
        self.connect()


if __name__ == '__main__':
    with DBManage() as db:
        print(db.show_data('your_table_name'))
